package io.tec.cloud.algorithm.c05_v20220617;

/**
 * https://leetcode.com/problems/count-of-range-sum/
 * <p>
 * 给定一个数组arr，两个整数lower和upper，
 * <p>
 * 返回arr中有多少个子数组的累加和在[lower,upper]范围上
 */
public class Code01_CountOfRangeSum {


//    public static void main(String[] args) {
//        int[] arr = {-2, 5, -1};
//        System.out.println(countRangeSum(arr, -2, 2));
//    }

    public  int countRangeSum(int[] arr, int lower, int upper) {
        if (arr == null || arr.length < 1) {
            return 0;
        }

        long[] sum = new long[arr.length];
        sum[0] = arr[0];
        for (int i = 1; i < arr.length; i++) {
            sum[i] = sum[i - 1] + arr[i];
        }

        return process(sum, 0, sum.length - 1, lower, upper);
    }

    public static int process(long[] sum, int l, int r, int lower, int upper) {
        // base case
        if (l == r) {
            if (sum[l] >= lower && sum[l] <= upper) {
                return 1;
            }
            return 0;
        }
        int mid = l + ((r - l) >> 1);


        int leftPart = process(sum, l, mid, lower, upper);
        int rightPart = process(sum, mid + 1, r, lower, upper);
        int merge = merge2(sum, l, mid, r, lower, upper);
        return leftPart + rightPart + merge;
    }

    public static int merge(int[] arr, int L, int M, int R, int lower, int upper) {
        int ans = 0;
        int windowL = L;
        int windowR = L;
        // [windowL, windowR)
        for (int i = M + 1; i <= R; i++) {
            long min = arr[i] - upper;
            long max = arr[i] - lower;
            while (windowR <= M && arr[windowR] <= max) {
                windowR++;
            }
            while (windowL <= M && arr[windowL] < min) {
                windowL++;
            }
            ans += windowR - windowL;
        }
        int[] help = new int[R - L + 1];
        int i = 0;
        int p1 = L;
        int p2 = M + 1;
        while (p1 <= M && p2 <= R) {
            help[i++] = arr[p1] <= arr[p2] ? arr[p1++] : arr[p2++];
        }
        while (p1 <= M) {
            help[i++] = arr[p1++];
        }
        while (p2 <= R) {
            help[i++] = arr[p2++];
        }
        for (i = 0; i < help.length; i++) {
            arr[L + i] = help[i];
        }
        return ans;
    }


    private static int merge2(long[] sum, int l, int mid, int r, int lower, int upper) {
        long[] help = new long[r - l + 1];
        int lowIndex = l;
        int upperIndex = l;
        int li = l;
        int ri = mid + 1;
        int count = 0;
        while (ri <= r) {
            // sum
            while (upperIndex <= mid && sum[upperIndex] <= sum[ri] - lower) {
                upperIndex++;
            }
            while (lowIndex <= mid && sum[lowIndex] < sum[ri] - upper) {
                lowIndex++;
            }
            count += (upperIndex - lowIndex);
            ri++;
        }
        li = l;
        ri = mid + 1;
        int idx = 0;
        while (li <= mid && ri <= r) {
            if (sum[li] <= sum[ri]) {
                help[idx++] = sum[li++];
            } else {
                help[idx++] = sum[ri++];
            }
        }

        while (li <= mid) {
            help[idx++] = sum[li++];
        }

        while (ri <= r) {
            help[idx++] = sum[ri++];
        }
        for (int i = 0; i < help.length; i++) {
            sum[l + i] = help[i];
        }
        return count;
    }
}
